import networkx as nx
import matplotlib.pyplot as plt
import numpy as np
import random
import math
from collections import defaultdict

# 1. 生成有向管网拓扑（确保流向高压→低压）
def generate_gas_network(num_nodes=50, num_sources=3, num_users=15, valve_prob=0.4):
    """生成带压力梯度的有向管网"""
    G = nx.DiGraph()
    
    # 生成节点位置（确保气源在中心区域）
    center = (0.5, 0.5)
    positions = {}
    for i in range(num_nodes):
        if i < num_sources:
            positions[i] = (
                center[0] + random.uniform(-0.2, 0.2),
                center[1] + random.uniform(-0.2, 0.2)
            )
        else:
            positions[i] = (random.random(), random.random())
    
    # 设置节点压力梯度
    for i in range(num_nodes):
        if i < num_sources:
            G.add_node(i, type='source', pressure=random.uniform(0.8, 1.2), pos=positions[i])
        elif i < num_sources + num_users:
            G.add_node(i, type='user', pressure=random.uniform(0.4, 0.6), 
                      priority=random.randint(1, 5), pressure_req=random.uniform(0.2, 0.4),
                      pos=positions[i])
        else:
            G.add_node(i, type='junction', pressure=random.uniform(0.5, 0.7), pos=positions[i])
    
    # 添加带流向的管道（高压→低压）
    for i in range(num_nodes):
        for j in range(i+1, num_nodes):
            if random.random() < 0.3 and math.dist(positions[i], positions[j]) < 0.4:
                # 根据压力确定流向
                if G.nodes[i]['pressure'] > G.nodes[j]['pressure']:
                    u, v = i, j
                else:
                    u, v = j, i
                
                # 添加管道属性
                diameter = random.choice([0.1, 0.15, 0.2, 0.25])
                length = math.dist(positions[i], positions[j]) * 10
                
                # 添加带方向的边
                G.add_edge(u, v, 
                          diameter=diameter,
                          length=length,
                          resistance=random.uniform(0.02, 0.05),
                          valve=random.random() < valve_prob)
                
                # 设置阀门属性
                if G.edges[u, v]['valve']:
                    G.edges[u, v]['valve_time'] = random.choice([1, 2, 3, 5, 10])
    
    return G, positions

# 2. 流向验证函数（防止逆流）
def validate_flow_direction(G, sources):
    """确保所有路径都是从气源向外辐射"""
    for node in G.nodes:
        if node in sources:
            continue
            
        # 检查是否存在指向气源的路径（禁止逆流）
        for source in sources:
            if G.has_edge(node, source):
                G.remove_edge(node, source)
                print(f"Removed reverse flow edge: {node}->{source}")
    
    return G

# 3. 泄漏点模拟
def identify_leaks(G, num_leaks=2):
    """在非气源节点生成泄漏点"""
    non_source_nodes = [n for n in G.nodes if G.nodes[n]['type'] != 'source']
    leak_nodes = random.sample(non_source_nodes, num_leaks)
    
    leaks = {}
    for node in leak_nodes:
        pressure = G.nodes[node]['pressure']
        leak_rate = random.uniform(0.8, 1.2) * pressure * math.sqrt(pressure)
        leaks[node] = (G.nodes[node]['pos'], leak_rate)
    
    return leaks

# 4. 最小割算法（有向图版本）[2,7](@ref)
def find_min_cut(G, leaks, sources):
    """基于流向的有向图最小割算法"""
    flow_graph = nx.DiGraph()
    super_source = "super_source"
    super_sink = "super_sink"
    flow_graph.add_node(super_source)
    flow_graph.add_node(super_sink)
    
    # 构建带容量的流向图
    for u, v in G.edges():
        if G.edges[u, v].get('valve', False):
            capacity = 1
            flow_graph.add_edge(u, v, capacity=capacity)
    
    # 连接超级源点->气源
    for source in sources:
        flow_graph.add_edge(super_source, source, capacity=10000)
    
    # 连接泄漏点->超级汇点
    for leak in leaks:
        flow_graph.add_edge(leak, super_sink, capacity=10000)
    
    # 计算最小割[2](@ref)
    try:
        min_cut_value, partition = nx.minimum_cut(flow_graph, super_source, super_sink)
        reachable, non_reachable = partition
        
        # 提取割边（实际阀门）
        min_cut_set = []
        for u, v in flow_graph.edges():
            if u in reachable and v in non_reachable:
                if G.has_edge(u, v) and G.edges[u, v].get('valve', False):
                    min_cut_set.append((u, v))
        
        # 识别被隔离的节点（泄漏点连通分量）
        H = G.copy()
        H.remove_edges_from(min_cut_set)
        isolated_nodes = set()
        for leak in leaks:
            if leak in H.nodes:
                isolated_nodes |= set(nx.descendants(H, leak)) | {leak}
                
        return min_cut_set, isolated_nodes
        
    except nx.NetworkXUnbounded:
        print("Error: Unbounded flow detected, check flow directions")
        return [], set()

# 5. 计算受影响用户
def get_affected_users(G, isolated_nodes, sources):
    """获取被隔离区域内的用户节点"""
    disconnected_users = []
    for node in isolated_nodes:
        if G.nodes[node]['type'] == 'user':
            # 检查用户是否仍然连接到气源
            connected = False
            for source in sources:
                if nx.has_path(G, source, node):
                    connected = True
                    break
            if not connected:
                disconnected_users.append(node)
    return disconnected_users

# 6. 管网可视化
def visualize_network(G, pos, valves_to_close=None, leaks=None, disconnected_users=None, sources=None):
    """可视化管网与关阀方案"""
    plt.figure(figsize=(15, 12))
    
    # 节点颜色分类
    node_colors = []
    for node in G.nodes:
        if leaks and node in leaks:
            node_colors.append('red')  # 泄漏点
        elif disconnected_users and node in disconnected_users:
            node_colors.append('purple')  # 断气用户
        elif sources and node in sources:
            node_colors.append('green')  # 气源
        elif G.nodes[node]['type'] == 'user':
            node_colors.append('blue')  # 正常用户
        else:
            node_colors.append('lightgray')  # 其他节点
    
    # 绘制节点
    nx.draw_networkx_nodes(G, pos, node_color=node_colors, node_size=300, edgecolors='black')
    
    # 绘制正常边（带箭头表示流向）
    nx.draw_networkx_edges(
        G, pos, 
        edgelist=[e for e in G.edges() if e not in (valves_to_close or [])],
        edge_color='gray',
        width=1.5,
        arrows=True,
        arrowstyle='->',
        arrowsize=15
    )
    
    # 绘制需关闭的阀门
    if valves_to_close:
        nx.draw_networkx_edges(
            G, pos, 
            edgelist=valves_to_close,
            edge_color='red',
            width=3,
            style='dashed',
            arrows=True,
            arrowstyle='->',
            arrowsize=20
        )
        # 标记阀门位置
        for u, v in valves_to_close:
            if u in pos and v in pos:
                mid_x = (pos[u][0] + pos[v][0]) / 2
                mid_y = (pos[u][1] + pos[v][1]) / 2
                plt.scatter(mid_x, mid_y, s=300, c='red', marker='X', zorder=3)
    
    # 图例
    plt.legend(handles=[
        plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='green', markersize=10, label='气源'),
        plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='blue', markersize=10, label='用户'),
        plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='red', markersize=10, label='泄漏点'),
        plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='purple', markersize=10, label='断气用户'),
        plt.Line2D([0], [0], color='gray', lw=2, label='正常管道'),
        plt.Line2D([0], [0], color='red', lw=2, linestyle='dashed', label='需关阀门')
    ], loc='upper right', fontsize=12)
    
    plt.title("Gas Network Isolation - Min Cut Solution", fontsize=14)
    plt.axis('off')
    plt.tight_layout()
    plt.savefig("gas_network_min_cut.png", dpi=300)
    plt.show()

# 7. 主流程
def main():
    print("=== 燃气泄漏应急关阀优化系统（最小割算法） ===")
    
    print("1. 生成带压力梯度的有向管网...")
    G, pos = generate_gas_network(num_nodes=40, num_sources=3, num_users=12)
    
    sources = [n for n, data in G.nodes(data=True) if data['type'] == 'source']
    print("2. 验证流向（防止逆流）...")
    G = validate_flow_direction(G, sources)
    
    print("3. 模拟泄漏点...")
    leaks = identify_leaks(G, num_leaks=2)
    leak_nodes = list(leaks.keys())
    print(f"  泄漏点位置: {leak_nodes}")
    
    print("4. 查找最小阀门集合隔离泄漏点...")
    min_cut_set, isolated_nodes = find_min_cut(G, leak_nodes, sources)
    print(f"  需关闭的阀门数: {len(min_cut_set)}")
    
    print("5. 计算受影响用户...")
    disconnected_users = get_affected_users(G, isolated_nodes, sources)
    print(f"  受影响用户数: {len(disconnected_users)}")
    
    # 计算最大关阀时间和总泄漏量
    max_time = max([G.edges[v]['valve_time'] for v in min_cut_set]) if min_cut_set else 0
    total_leak = sum(leak_data[1] for leak_data in leaks.values()) * max_time
    print(f"  最大关阀时间: {max_time:.1f} 分钟")
    print(f"  预估总泄漏量: {total_leak:.1f} 单位")
    
    print("6. 可视化关阀方案...")
    visualize_network(G, pos, min_cut_set, leak_nodes, disconnected_users, sources)
    
    print("=== 处理完成 ===")

if __name__ == "__main__":
    main()